import os
import torch

# 保存模型
def save(model, optimizer, epoch, step, error, best_error):
    path = os.path.join('checkpoints', '{:03}epoch.pth'.format(epoch))

    # 保存模型权重到字典
    state = {}
    state['state_dict'] = model.state_dict()
    state['optimizer'] = optimizer.state_dict()
    state['epoch'] = epoch
    state['step'] = step
    state['vali_avg_error'] = error['avg_error']
    state['vali_error_1px'] = error['error_1px']
    state['vali_error_2px'] = error['error_2px']
    state['vali_error_3px'] = error['error_3px']
    state['vali_error_4px'] = error['error_4px']
    state['vali_error_5px'] = error['error_5px']

    torch.save(state, path)

    # 保存最低平均误差的权重
    avg_error = error['avg_error']
    if avg_error < best_error:
        best_error = avg_error

        path = os.path.join('checkpoints', 'best.pth'.format(epoch))

        # 保存模型权重到字典
        state = {}
        state['state_dict'] = model.state_dict()
        state['optimizer'] = optimizer.state_dict()
        state['epoch'] = epoch
        state['step'] = step
        state['vali_avg_error'] = error['avg_error']
        state['vali_error_1px'] = error['error_1px']
        state['vali_error_2px'] = error['error_2px']
        state['vali_error_3px'] = error['error_3px']
        state['vali_error_4px'] = error['error_4px']
        state['vali_error_5px'] = error['error_5px']

        torch.save(state, path)

        print('best model in step {:05} epoch {:03}'.format(step, epoch))

    return best_error
